from pypinyin import lazy_pinyin   #实现汉字转拼音的效果
import pickle   #加载pkl文件
import os
import faiss    #使用的是faiss向量数据库
from emb_model import MyEmbModel  #这里使用的是m3e-base词向量
from load_documents import MyDocument   #进行文档加载


def cover_name(name):
    #例如给定一个”你好“,则返回   ni_hao
    return "_".join(lazy_pinyin(name))


class MyEmbDatabase():
    def __init__(self, emb_dir, qa_df, name):
        #声明client为一个全局变量
        global client
        self.name = cover_name(name)
        self.emb_model = MyEmbModel(emb_dir) #定义自己的向量模型

        #如果不存在question向量化后的pkl
        if not os.path.exists(os.path.join(".cache", f"{name}_faiss_index.pkl")):
            #采用IndexFlatL2方式，欧式距离，传入所用向量库的维度  ，构建faiss索引
            index = faiss.IndexFlatL2(self.emb_model.model.get_sentence_embedding_dimension())
            #将所有提取出来的qa问答对的question进行向量化
            embs = self.emb_model.to_emb(qa_df["question"])
            #faiss索引中添加这些question的向量
            index.add(embs)

            #将question索引写入faiss_index.pkl
            with open(os.path.join(".cache", f"{name}_faiss_index.pkl"), "wb") as f:
                pickle.dump(index, f)
            #将question and answer 写入faiss_index_df.pkl中
            with open(os.path.join(".cache", f"{name}_faiss_index_df.pkl"), "wb") as f:
                pickle.dump(qa_df, f)
        else:
            with open(os.path.join(".cache", f"{name}_faiss_index.pkl"), "rb") as f:
                index = pickle.load(f)
            with open(os.path.join(".cache", f"{name}_faiss_index_df.pkl"), "rb") as f:
                qa_df = pickle.load(f)
        self.index = index    #index为qestion所有向量化
        self.qa_df = qa_df    #qa_df为所有question-answer向量化

    def search(self, content, topn=3):
        global client
        #content进行向量化
        if isinstance(content, str):
            content = self.emb_model.to_emb(content)
        #返回与该content最相关的topn个
        distances, idxs = self.index.search(content,topn)
        #将最相关的结果返回  idxs[0] 是一个[]  里面是最相关的3个的排序，第一个表示最相关的，距离最短的
        results = self.qa_df.iloc[idxs[0]]
        #这里resluts返回的是topn个相关的q-a问答对
        return results


class MyDataBase:
    def __init__(self, path, name=None):
        # #没有name的化name就是路径
        # if name is None:
        #     name = path
        # print(os.path.join(path, "txt"))

        #查看有没有对应数据库
        if not os.path.exists(os.path.join(path, "txt")):
            print("对应的知识库txt文件夹不存在")
            exit(-999)
        #查看有没有对对应切片文档进行qa对抽取
        if not os.path.exists(os.path.join(path, "qa")):
            print("对应的知识库没有进行QA问答对抽取")
            exit(-999)
        #没有.cache文件就创建一个chche文件
        if not os.path.exists(".cache"):
            os.mkdir(".cache")
        #前提是1.需要有知识库，2.对文档进行抽取，3.有.cache文件夹
        self.name = name   #该知识库的名称
        self.document = MyDocument(os.path.join(path, "txt"), name)  #该知识库的文档
        self.emb_database = None   #该知识库的向量数据库

    #创建向量数据库
    def create_emb_database(self, qa_df):
        self.emb_database = MyEmbDatabase("moka-ai_m3e-base", qa_df, self.name)
        self.qa_df = qa_df

    #通过qa问答对找到最相关的三个qa问答对
    def search(self, text, topn=3):
        if not self.emb_database:
            self.emb_database = MyEmbDatabase("moka-ai_m3e-base", None, self.name)
        return self.emb_database.search(text, topn)


def load_database(dir_path="data/database_dir"):
    #列出所有文件夹['小米汽车']
    '''
    这里相当于
    for(){
        if()
    }
    首先检查data/database_dir下的所有文件夹（相当于一个知识库），如果该文件夹是一个文件夹的情况下，再将该文件夹存到dirs中
    '''
    dirs = [name for name in os.listdir(dir_path) if os.path.isdir(f"{dir_path}/{name}")]

    database_list = []  #添加的是MydataBase对象
    database_namelist = []  #添加的是所有数据库名称

    for dir in dirs:
        database = MyDataBase(f"{dir_path}/{dir}", dir)

        database_list.append(database)
        database_namelist.append(dir)

    return database_list, database_namelist




